

import codecs

# document states
PRE_DOCUMENT, IN_ELEMENT, IN_CONTENT, END_DOCUMENT = range(4)

entities = [('&','amp'),
            ('<', 'lt'),
            ('>', 'gt')]
quotes = [('"', 'quot'),
          ("'", 'apos')]

unisome = [(unicode(x), unicode('&%s;' % y)) for x, y in entities]
strsome = [(x, '&%s;' % y) for x, y in entities]
uniall = [(unicode(x), unicode('&%s;' % y)) for x, y in entities + quotes]
strall = [(x, '&%s;' % y) for x, y in entities + quotes]

def _getmap(isuni, quotes):
    if isuni:
        if quotes:
            return uniall
        return unisome
    if quotes:
        return strall
    return strsome    

def _escape(text, quotes = False):
    map = _getmap(type(text) is unicode, quotes)
    for (src, dest) in map:
        text = text.replace(src, dest)
    return text

class CurrentNamespace:
    pass

class XmlWriter:
    '''An xml output class, writes to an arbitrary buffer.  Based loosely
    on a similar C++ implementation available here:
    http://www.deez.info/sengelha/blog/2006/01/29/implementing-ixmlwriter-series/
    '''
    def __init__(self, outp = None, encoding = 'utf-8', pretty = True):
        self.encoding = encoding
        enc, dec, read_wrap, write_wrap = codecs.lookup(encoding)
        if outp is None:
            from cStringIO import StringIO
            self.buf = StringIO()
            self._out = self.buf
        else:
            self._out = outp
            self.buf = None
        self._out = write_wrap(self._out)
        self._write = self._out.write
        self.pretty = pretty
        self.state = PRE_DOCUMENT
        self.stack = []
        self.fwdnsmap = {}
        self.backnsmap = {}
        self._write(u'<?xml version="1.0" encoding="%s" ?>\n' % encoding)

    def _indent(self, text, extra = False):
        if self.pretty:
            n = len(self.stack)
            if extra: n += 1
            self._write(u'  ' * n)
        self._write(text)

    def _endline(self):
        if self.pretty:
            self._write(u'\n')

    def _checkStartEl(self):
        if self.state == IN_ELEMENT:
            self._write(u'>')
            self._endline()
            self.state = IN_CONTENT

    def _remap(self):
        fwd = self.fwdnsmap
        fwd.clear()
        back = self.backnsmap
        back.clear()
        for n in range(len(self.stack)):
            name, ns, prefix, mappings = self.stack[-n]
            if ns:
                if not prefix in back:
                    back[prefix] = ns
                if not ns in fwd:
                    fwd[ns] = prefix
            for prefix, ns in mappings:
                if not prefix in back:
                    back[prefix] = ns
                if not ns in fwd:
                    fwd[ns] = prefix

    def _findprefix(self, ns, invent = False):
        if ns in ('xmlns', u'xmlns', 'xml', u'xml'):
            return ns
        if ns in self.fwdnsmap:
            return self.fwdnsmap[ns]
        if invent:
            for n in range(50):
                pref = 'ns%d' % n
                if not pref in fwd:
                    self.bindNS(pref, ns)
                    return ns

    def comment(self, text):
        self._checkStartEl()
        if u'--' in text:
            raise Exception("Invalid comment: connot contain '--': %s" % text)
        self._indent(u"<!-- %s -->" % text)
        self._endline()

    def contentElement(self, name, content, ns = None, nsprefix = None):
        self.startElement(name, ns, nsprefix)
        self.content(content)
        self.endElement()
                
    def startElement(self, name, ns = None, nsprefix = None):
        if self.state == END_DOCUMENT:
            raise Exception("No elements allowed after root is closed")
        self._checkStartEl()
        prefix = nsprefix
        if ns and not prefix:
            prefix = self._findprefix(ns)
        nses = [x[1] for x in self.stack if x[1] and not x[2]]
        defaultns = None
        if nses and (nses[-1] == ns):
            defaultns = ns
        if prefix and not defaultns:
            self._indent(u'<%s:%s' % (prefix, name))
        else:
            self._indent('<%s' % name)
        self.state = IN_ELEMENT
        self.attrs = 0
        if not prefix and not defaultns:
            self.attribute(u"xmlns", ns)
        self.stack.append((name, ns, nsprefix, ()))
        if ns and nsprefix:
            self.bindNS(nsprefix, ns)
        self._remap()

    def attribute(self, name, value, ns = None):
        if self.state == IN_ELEMENT:
            prefix = None
            if ns: prefix = self._findprefix(ns)
            if self.attrs and self.pretty:
                self._endline()
                self._indent(u'  ')
            if not prefix:
                self._write(u' %s="%s"' % (name, _escape(value, quotes=1)))
            else:
                self._write(u' %s:%s="%s"' % (prefix, name, value))
            self.attrs = 1
        else:
            raise Exception("Can't write attribute when not in element")

    def attributes(self, attrs = None, **kw):
        if type(attrs) is dict:
            for k in attrs.keys():
                self.attribute(k, unicode(attrs[k]))
        for k in kw.keys():
            self.attribute(k, unicode(kw[k]))

    def bindNS(self, newprefix, newns):
        if self.state == IN_ELEMENT:
            tup = ((newprefix, newns),)
            name, ns, prefix, prefmap = self.stack.pop()
            prefmap += tup
            self.stack.append((name, ns, prefix, prefmap))
            self.attribute(newprefix, newns, u'xmlns')
            self._remap()
        else:
            raise Exception("Can't bind namespace except inside element start")
        
    def endElement(self):
        (name, ns, prefix, nsmap) = self.stack.pop()
        if ns or prefix or nsmap:
            self._remap()
        if self.state == IN_ELEMENT:
            self._write(u' />')
        else:
            self._indent('</%s>' % name)
        if self.stack:
            self.state = IN_CONTENT
        else:
            self.state = END_DOCUMENT
        self._endline()

    def content(self, text):
        self._checkStartEl()
        if self.state == IN_CONTENT:
            self._indent(_escape(text))
            self._endline()
        else:
            raise Exception("Can't add text content before or after root element")

    def close(self):
        while self.stack:
            self.endElement()
        self._out.flush()

if __name__ == '__main__':
    import sys
    w = XmlWriter(outp = sys.stdout, pretty = True)
    w.comment(u"this is a test")
    w.startElement(u"abc", u'http://meat.net/svg/')
    w.bindNS(u"xlink", u"http://www.w3.org/1999/xlink")
    w.startElement(u"def", u'http://meat.net/test/')
    w.attribute(u'href', u'http://meat.net/', u'http://www.w3.org/1999/xlink')
    w.endElement()
    w.startElement(u"ghi", u'http://meat.net/svg/')
    w.attribute(u"jkl", u'<&> is illegal')
    w.content(u"this is ghi content, and we're testing >escaping & stuff< ")
    w.close()
